#include <stdio.h>
#include <math.h>
#include <stdint.h>
#include <string.h>
#include <cuda_runtime.h>
#include "cuda_cal.h"

// 错误检查宏
#define CHECK(call) \
    do { \
        cudaError_t err = call; \
        if (err != cudaSuccess) { \
            fprintf(stderr, "Error: %s in file %s, line %d\n", \
                    cudaGetErrorString(err), __FILE__, __LINE__); \
            exit(EXIT_FAILURE); \
        } \
    } while (0)

#define CHECK_RETURN_NULL(call) \
    do { \
        cudaError_t err = call; \
        if (err != cudaSuccess) { \
            fprintf(stderr, "Error: %s in file %s, line %d\n", \
                    cudaGetErrorString(err), __FILE__, __LINE__); \
            return nullptr; \
        } \
    } while (0)

// 1. CPU内存
static PointXYZI_Float* h_points;
static double* h_azimuth;
static double* h_pitch;
static double* h_distance;
static PointResult point_result;

// 2. GPU内存
static PointXYZI_Float* d_points;
static double* d_azimuth;
static double* d_pitch;
static double* d_distance;

// 3.总个数
static uint32_t total;

// CUDA核函数：并行计算每个点的azimuth, pitch, distance
__global__ void gpu_cal(const PointXYZI_Float* points, double* azimuth, double* pitch, double* distance, int total) {
    // 计算当前线程索引
    int idx = blockIdx.x * blockDim.x + threadIdx.x;

    // 确保不越界访问
    if (idx < total) {
        const PointXYZI_Float& p = points[idx];

        // 使用CUDA设备端数学函数
        azimuth[idx] = atan2f(p.y, p.x) * 57.2958;
        float r = p.x * p.x + p.y * p.y;
        pitch[idx] = atan2f(p.z, sqrtf(r)) * 57.2958;
        distance[idx] = sqrtf(r + p.z * p.z);
    }
}

__global__ void gpu_cal_v2(const PointXYZI_Float* points, double* azimuth, double* pitch, double* distance, int total) {
    // 计算当前线程索引
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int stride = blockDim.x * gridDim.x;

    // 确保不越界访问
    for (int i = idx; i < total; i += stride)
    {
        const PointXYZI_Float& p = points[i];

        // 使用CUDA设备端数学函数
        azimuth[i] = atan2f(p.y, p.x) * 57.2958;
        float r = p.x * p.x + p.y * p.y;
        pitch[i] = atan2f(p.z, sqrtf(r)) * 57.2958;
        distance[i] = sqrtf(r + p.z * p.z);
    }
}

bool init_gpu_mem()
{
    const uint32_t width = 1800;
    const uint32_t height = 32 * 2;
    total = width * height;

    // 1. 分配Pinned CPU内存（页锁定内存） 
    CHECK(cudaMallocHost(&h_points, total * sizeof(struct PointXYZI_Float)));
    CHECK(cudaMallocHost(&h_azimuth, total * sizeof(double)));
    CHECK(cudaMallocHost(&h_pitch, total * sizeof(double)));
    CHECK(cudaMallocHost(&h_distance, total * sizeof(double)));

    // 2. 分配gpu 内存
    CHECK(cudaMalloc(&d_points, total * sizeof(struct PointXYZI_Float)));
    CHECK(cudaMalloc(&d_azimuth, total * sizeof(double)));
    CHECK(cudaMalloc(&d_pitch, total * sizeof(double)));
    CHECK(cudaMalloc(&d_distance, total * sizeof(double)));

    if(!h_points || !h_azimuth || !h_pitch || !h_distance || 
        !d_points || !d_azimuth || !d_pitch || !d_distance)
    {
        printf("cuda memory malloc failed\n");
        return false;
    }

    return true;
}

void release_gpu_mem()
{
    // 1.释放gpu内存
    CHECK(cudaFree(d_points));
    CHECK(cudaFree(d_azimuth));
    CHECK(cudaFree(d_pitch));
    CHECK(cudaFree(d_distance));

    // 2.释放cpu内存
    CHECK(cudaFreeHost(h_points));    // 使用cudaFreeHost释放pinned内存
    CHECK(cudaFreeHost(h_azimuth));
    CHECK(cudaFreeHost(h_pitch));
    CHECK(cudaFreeHost(h_distance));
}

PointResult* gpu_cal_atan(PointXYZI_Float* host_points)
{
    if (host_points == nullptr)
        return nullptr;

    // 1. 将数据从CPU（pinned内存）复制到GPU
    // if(host_points != h_points)
    //     memcpy(h_points, host_points, total * sizeof(PointXYZI_Float));
    CHECK_RETURN_NULL(cudaMemcpy(d_points, host_points, total * sizeof(PointXYZI_Float), cudaMemcpyHostToDevice));

    // 2. 配置CUDA核函数并执行
    dim3 blockDim(256);
    dim3 gridDim((total + blockDim.x - 1) / blockDim.x);
    gpu_cal <<<gridDim, blockDim>>> (d_points, d_azimuth, d_pitch, d_distance, total);

    // int blockSize = 256;
    // int blockNum = (total + blockSize - 1) / blockSize;
    // gpu_cal_v2 <<<blockNum, blockSize>>> (d_points, d_azimuth, d_pitch, d_distance, total);

    CHECK_RETURN_NULL(cudaGetLastError());
    CHECK_RETURN_NULL(cudaDeviceSynchronize());

    // 3. 将结果从GPU复制回CPU（pinned内存）
    CHECK_RETURN_NULL(cudaMemcpy(h_azimuth, d_azimuth, total * sizeof(double), cudaMemcpyDeviceToHost));
    CHECK_RETURN_NULL(cudaMemcpy(h_pitch, d_pitch, total * sizeof(double), cudaMemcpyDeviceToHost));
    CHECK_RETURN_NULL(cudaMemcpy(h_distance, d_distance, total * sizeof(double), cudaMemcpyDeviceToHost));

    // 4. 返回结果
    point_result.azimuth = h_azimuth;
    point_result.pitch = h_pitch;
    point_result.distance = h_distance;
    
    return &point_result;
}

void test()
{
    if(!init_gpu_mem())
    {
        fprintf(stderr, "GPU内存初始化失败\n");
        return;
    }

    // 初始化点数据
    for (uint32_t i = 0; i < total; i++) {
        h_points[i].x = 1.0f;
        h_points[i].y = 2.0f;
        h_points[i].z = 3.0f;
    }
    PointResult* result = gpu_cal_atan(h_points);

    // 验证结果
    printf("第一个点的计算结果:\n");
    printf("azimuth: %f\n", result->azimuth[0]);
    printf("pitch: %f\n", result->pitch[0]);
    printf("distance: %f\n", result->distance[0]);

    release_gpu_mem();
}


int main()
{
    test();
    return 0;
}
